from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client
from langchain_mcp_adapters.tools import load_mcp_tools
from langgraph.prebuilt import create_react_agent
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
import asyncio
import sys
import os

# 添加当前目录到 Python 路径
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

load_dotenv()
model = ChatOpenAI(model="deepseek-chat")

async def main():
    # 获取当前 Python 解释器路径
    python_executable = sys.executable
    
    # 获取当前环境的 PYTHONPATH
    current_pythonpath = os.environ.get("PYTHONPATH", "")
    
    # 添加当前目录到 PYTHONPATH
    new_pythonpath = os.path.dirname(os.path.abspath(__file__))
    if current_pythonpath:
        new_pythonpath = new_pythonpath + os.pathsep + current_pythonpath
    
    # 复制当前环境变量并修改 PYTHONPATH
    env = os.environ.copy()
    env["PYTHONPATH"] = new_pythonpath
    
    server_params = StdioServerParameters(
        command=python_executable,  # 使用当前 Python 解释器
        args=["D:\\my\\python\\mcp\\math_server.py"],
        env=env  # 使用修改后的环境变量
    )

    async with stdio_client(server_params) as (read, write):
        async with ClientSession(read, write) as session:
            # Initialize the connection
            await session.initialize()

            # Get tools
            tools = await load_mcp_tools(session)

            # Create and run the agent
            agent = create_react_agent(model, tools)
            agent_response = await agent.ainvoke({"messages": [{"role": "user", "content": "what's (3 + 5) x 12?"}]})
            
            # 提取计算结果
            if "messages" in agent_response:
                messages = agent_response["messages"]
                # 查找工具消息，特别是最后一个工具消息的结果
                tool_messages = [msg for msg in messages if hasattr(msg, 'name') and hasattr(msg, 'content')]
                
                if tool_messages:
                    # 获取最后一个工具消息的内容作为最终结果
                    final_result = tool_messages[-1].content
                    print(f"计算结果: {final_result}")
                    
                    # 打印计算过程
                    print("\n计算过程:")
                    for msg in tool_messages:
                        print(f"- {msg.name}: {msg.content}")
                else:
                    print("未找到工具消息")
            else:
                print("响应中没有 messages 键")

# 运行异步主函数
if __name__ == "__main__":
    asyncio.run(main())